#include <windows.h>
#include <iostream>
#include "opencv2/myopencv.hpp"
#include "NuiApi.h"
#pragma comment(lib, "Kinect10.lib")

using namespace std;

class KinectDepthWrapper {
private:
	INuiSensor *m_pNuiSensor; 
	HANDLE m_pDepthStreamHandle;
	HANDLE m_hNextDepthFrameEvent;

public:
	int width, height;

	void init(int w, int h) {
		HRESULT hr = NuiCreateSensorByIndex(0, &m_pNuiSensor);
		if (FAILED(hr)) cout << "create sensor failed." << endl;

		WCHAR *id = m_pNuiSensor->NuiDeviceConnectionId();
		m_hNextDepthFrameEvent = CreateEvent(NULL, TRUE, FALSE, NULL);

		hr = m_pNuiSensor->NuiInitialize(NUI_INITIALIZE_FLAG_USES_DEPTH);
		if (FAILED(hr)) cout << "failed init." << endl;

		NUI_IMAGE_RESOLUTION resolution = NUI_IMAGE_RESOLUTION_640x480;
		switch (w) {
		case 80: resolution = NUI_IMAGE_RESOLUTION_80x60; break;
		case 320: resolution = NUI_IMAGE_RESOLUTION_320x240; break;
		case 640: resolution = NUI_IMAGE_RESOLUTION_640x480; break;
		case 1280: resolution = NUI_IMAGE_RESOLUTION_1280x960; break;
		default:cout << "unsupport resolution. use 640*480." << endl;
		}

		hr = m_pNuiSensor->NuiImageStreamOpen(
			NUI_IMAGE_TYPE_DEPTH,
			resolution,
			0,
			2,
			m_hNextDepthFrameEvent,
			&m_pDepthStreamHandle);

		width = w;
		height = h;

		//user near mode
		NuiImageStreamSetImageFrameFlags(m_pDepthStreamHandle, NUI_IMAGE_STREAM_FLAG_ENABLE_NEAR_MODE);
	}

	bool retrieve(short *sPtr) {
		DWORD retVal;
		DWORD frameWidth, frameHeight;

		retVal = WaitForSingleObject(m_hNextDepthFrameEvent, 100);
		if (retVal != WAIT_OBJECT_0) return false;

		NUI_IMAGE_FRAME imageFrame;
		HRESULT hr = m_pNuiSensor->NuiImageStreamGetNextFrame(
			m_pDepthStreamHandle,
			0,
			&imageFrame);
		INuiFrameTexture *pTexture = imageFrame.pFrameTexture;

		NUI_LOCKED_RECT lockedRect;
		pTexture->LockRect(0, &lockedRect, NULL, 0);
		if ( 0 != lockedRect.Pitch) {
			NuiImageResolutionToSize( imageFrame.eResolution, frameWidth, frameHeight );
			CV_Assert(frameHeight == height && frameWidth == width);
			memcpy(sPtr, lockedRect.pBits, sizeof(short)*height*width);
		}
		pTexture->UnlockRect(0);

		m_pNuiSensor->NuiImageStreamReleaseFrame( m_pDepthStreamHandle, &imageFrame );
		return true;
	}

	bool retrieve(CvMat *&ret) {
		if (ret == NULL) ret = cvCreateMat(height, width, CV_16SC1);
		return retrieve(ret->data.s);
	}
};

class KinectDepthColorWrapper {
private:
	INuiSensor *m_pNuiSensor; 
	HANDLE m_pDepthStreamHandle;
	HANDLE m_hNextDepthFrameEvent;
	HANDLE m_pColorStreamHandle;
	HANDLE m_hNextColorFrameEvent;

public:
	int width, height;
	long long colorT, depthT;
	long *coordinateBuff;

	~KinectDepthColorWrapper() {
		if (coordinateBuff != NULL)
			delete[] coordinateBuff;
	}

	void init(int w, int h) {
		coordinateBuff = NULL;
		HRESULT hr = NuiCreateSensorByIndex(0, &m_pNuiSensor);
		if (FAILED(hr)) cout << "create sensor failed." << endl;

		WCHAR *id = m_pNuiSensor->NuiDeviceConnectionId();
		m_hNextDepthFrameEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
		m_hNextColorFrameEvent = CreateEvent(NULL, TRUE, FALSE, NULL);

		hr = m_pNuiSensor->NuiInitialize(NUI_INITIALIZE_FLAG_USES_DEPTH | NUI_INITIALIZE_FLAG_USES_COLOR);
		if (FAILED(hr)) cout << "failed init." << endl;

		NUI_IMAGE_RESOLUTION resolution = NUI_IMAGE_RESOLUTION_640x480;
		switch (w) {
		case 80: resolution = NUI_IMAGE_RESOLUTION_80x60; break;
		case 320: resolution = NUI_IMAGE_RESOLUTION_320x240; break;
		case 640: resolution = NUI_IMAGE_RESOLUTION_640x480; break;
		case 1280: resolution = NUI_IMAGE_RESOLUTION_1280x960; break;
		default:cout << "unsupport resolution. use 640*480." << endl;
		}

		hr = m_pNuiSensor->NuiImageStreamOpen(
			NUI_IMAGE_TYPE_DEPTH,
			resolution,
			0,
			2,
			m_hNextDepthFrameEvent,
			&m_pDepthStreamHandle);
		hr = m_pNuiSensor->NuiImageStreamOpen(
			NUI_IMAGE_TYPE_COLOR,
			resolution,
			0,
			2,
			m_hNextColorFrameEvent,
			&m_pColorStreamHandle);

		width = w;
		height = h;

		//user near mode
		NuiImageStreamSetImageFrameFlags(m_pDepthStreamHandle, NUI_IMAGE_STREAM_FLAG_ENABLE_NEAR_MODE);
	}

	bool retrieve(short *sPtr, uchar *uPtr, long *colorCoordinates = NULL) { 
		DWORD retVal;
		DWORD frameWidth, frameHeight;
		HANDLE handles[] = {m_hNextColorFrameEvent, m_hNextDepthFrameEvent};

		retVal = WaitForMultipleObjects(2, handles, TRUE, 1000);

		if (retVal != WAIT_OBJECT_0) return false;

		NUI_IMAGE_FRAME depthFrame;
		NUI_IMAGE_FRAME colorFrame;

		HRESULT hr = m_pNuiSensor->NuiImageStreamGetNextFrame(
			m_pDepthStreamHandle,
			0,
			&depthFrame);
		hr = m_pNuiSensor->NuiImageStreamGetNextFrame(
			m_pColorStreamHandle,
			0,
			&colorFrame);

		INuiFrameTexture *pTexture = depthFrame.pFrameTexture;

		NUI_LOCKED_RECT lockedRect;
		pTexture->LockRect(0, &lockedRect, NULL, 0);
		if ( 0 != lockedRect.Pitch) {
			NuiImageResolutionToSize( depthFrame.eResolution, frameWidth, frameHeight );
			CV_Assert(frameHeight == height && frameWidth == width);
			memcpy(sPtr, lockedRect.pBits, sizeof(short)*height*width);
		}
		pTexture->UnlockRect(0);
		depthT = depthFrame.liTimeStamp.QuadPart;


		pTexture = colorFrame.pFrameTexture;
		pTexture->LockRect(0, &lockedRect, NULL, 0);
		if (0 != lockedRect.Pitch) {
			NuiImageResolutionToSize(colorFrame.eResolution, frameWidth, frameHeight);
			CV_Assert(frameHeight == height && frameWidth == width);
			byte *lrptr = lockedRect.pBits;
			for (int i=0; i<frameHeight; ++i)
				for (int j=0; j<frameWidth; ++j, uPtr+=3, lrptr+=4)
					memcpy(uPtr, lrptr, 3*sizeof(char));
		}
		pTexture->UnlockRect(0);
		colorT = colorFrame.liTimeStamp.QuadPart;

		if (colorCoordinates != NULL) {
			m_pNuiSensor->NuiImageGetColorPixelCoordinateFrameFromDepthPixelFrameAtResolution(
				colorFrame.eResolution,
				depthFrame.eResolution,
				width*height,
				(unsigned short *)sPtr,
				width*height*2,
				colorCoordinates
				);
		}

		m_pNuiSensor->NuiImageStreamReleaseFrame(m_pDepthStreamHandle, &depthFrame);
		m_pNuiSensor->NuiImageStreamReleaseFrame(m_pColorStreamHandle, &colorFrame);

		return true;
	}

	bool retrieve(CvMat *&depthRes, CvMat *&colorRes, long *colorCoordinates = NULL) {
		if (depthRes == NULL) depthRes = cvCreateMat(height, width, CV_16SC1);
		if (colorRes == NULL) colorRes = cvCreateMat(height, width, CV_8UC3);
		return retrieve(depthRes->data.s, colorRes->data.ptr, colorCoordinates);
	}

	bool retrieve(CvMat *&depthRes, CvMat *&colorRes, CvMat *&dcMapped) {
		if (coordinateBuff == NULL) coordinateBuff = new long[2*width*height];
		if (depthRes == NULL) depthRes = cvCreateMat(height, width, CV_16SC1);
		if (colorRes == NULL) colorRes = cvCreateMat(height, width, CV_8UC3);
		if (dcMapped == NULL) dcMapped = cvCreateMat(height, width, CV_16SC1);
		
		for (int i=0; i<height; ++i) {
			for (int j=0; j<width; ++j) {
				coordinateBuff[2*(i*width+j)]=j;
				coordinateBuff[2*(i*width+j)+1]=i;
			}
		}
		cvZero(dcMapped);

		if (retrieve(depthRes->data.s, colorRes->data.ptr, coordinateBuff)) {
			short *depthPtr = depthRes->data.s;
			for (int i=0; i<height; ++i) {
				for (int j=0; j<width; ++j, ++depthPtr) {
					int x = coordinateBuff[2*(i*width+j)];
					int y = coordinateBuff[2*(i*width+j)+1];
					short dval = *depthPtr;
					if (0<=x && x<width && 0<=y && y<height && dval > 100)
						cvSetReal2D(dcMapped, y,x, dval);
				}
			}
			return true;
		}
		else return false;
	}

};
